XGBoost Modeling R Notebook

LOAD DATA

df <- readRDS("../data/models/social-risk-crash-rate-data.rds")

HOT-ENCODE YEAR

df$year <- as.factor(df$year)
year_dummies <- model.matrix(~ year - 1, data = df)

df <- cbind(df[ , !(names(df) %in% c("year"))], year_dummies)

SELECT TRANSFORMATION FOR TARGET

compare_transformations <- function(df, target_var, features) {
  
  # Load required packages
  if (!require("MASS")) install.packages("MASS")
  if (!require("bestNormalize")) install.packages("bestNormalize")
  library(MASS)
  library(bestNormalize)
  
  y_raw <- df[[target_var]]
  X <- df[, features, drop = FALSE]
  X[] <- lapply(X, function(col) as.numeric(as.character(col)))
  
  # Log Transformation 
  y_log <- log1p(y_raw)
  lm_log <- lm(y_log ~ ., data = data.frame(y_log, X))
  pred_log <- predict(lm_log, newdata = data.frame(X))
  rmse_log <- sqrt(mean((y_log - pred_log)^2))
  
  # Box-Cox Transformation 
  y_bc_input <- ifelse(y_raw <= 0, 0.000001, y_raw)
  
  bc <- MASS::boxcox(y_bc_input ~ ., data = data.frame(y_bc_input, X), lambda = seq(-2, 2, 0.1))
  best_lambda <- bc$x[which.max(bc$y)]
  y_bc <- if(best_lambda == 0) {
    log(y_bc_input)
  } else {
    (y_bc_input^best_lambda - 1)/best_lambda
  }
  
  lm_bc <- lm(y_bc ~ ., data = data.frame(y_bc, X))
  pred_bc <- predict(lm_bc, newdata = data.frame(X))
  rmse_bc <- sqrt(mean((y_bc - pred_bc)^2))
  
  # Yeo-Johnson Transformation
  yj_obj <- yeojohnson(y_raw)
  y_yj <- predict(yj_obj)
  
  lm_yj <- lm(y_yj ~ ., data = data.frame(y_yj, X))
  pred_yj <- predict(lm_yj, newdata = data.frame(X))
  rmse_yj <- sqrt(mean((y_yj - pred_yj)^2))
  
  results <- data.frame(
    Transformation = c("Log (log1p)", "Box-Cox", "Yeo-Johnson"),
    RMSE = c(rmse_log, rmse_bc, rmse_yj),
    Best_Lambda_BoxCox = c(NA, best_lambda, NA)
  )
  
  print(results)
  
  return(list(
    log_model = lm_log,
    boxcox_model = lm_bc,
    yj_model = lm_yj,
    boxcox_lambda = best_lambda,
    transformation_results = results
  ))
}

# Select target variable
target_var <- "crash_rate_per_1000"

df <- df %>% dplyr::select(-injury_rate_per_1000, -fatality_rate_per_1000,  -total_population, -geoid)

features <- setdiff(names(df), c("crash_rate_per_1000", "borough"))
comparison <- compare_transformations(df, "crash_rate_per_1000", features)

NA

LOG TRANSFORM TARGET

X <- df %>% dplyr::select(-target_var, -borough)

# Log-transform the target variable
y <- log1p(df[[target_var]])  

hist(y, breaks = 50, main = "Log-Transformed Crash Rate", xlab = "log1p(crash_rate_per_1000)")

HYPERPARAMETER TUNING with Optuna

## CONVERT TO DMATRIX
dtrain_all <- xgb.DMatrix(data = as.matrix(X), label = y)

# Define Spatial Folds 
boroughs <- unique(df$borough)
folds <- lapply(boroughs, function(b) which(df$borough != b))
names(folds) <- boroughs

## Start Python venv
reticulate::use_virtualenv("r-reticulate", required = TRUE)

## OPTUNA-BASED SPATIAL CV
optuna <- import("optuna")

# Objective Function 
objective <- function(trial) {
  params <- list(
    booster = "gbtree",
    objective = "reg:squarederror",
    eval_metric = "rmse",
    eta = trial$suggest_float("eta", 0.01, 0.3, log = TRUE),
    max_depth = trial$suggest_int("max_depth", 3, 12),
    min_child_weight = trial$suggest_int("min_child_weight", 1, 10),
    subsample = trial$suggest_float("subsample", 0.5, 1.0),
    colsample_bytree = trial$suggest_float("colsample_bytree", 0.5, 1.0),
    gamma = trial$suggest_float("gamma", 0, 10),
    lambda = trial$suggest_float("lambda", 0, 10),
    alpha = trial$suggest_float("alpha", 0, 10)
  )
  
  cv_results <- xgb.cv(
    params = params,
    data = dtrain_all,
    nrounds = 150,
    folds = folds,
    early_stopping_rounds = 20,
    verbose = 0,
    nthread = detectCores() - 1
  )
  
  return(cv_results$evaluation_log$test_rmse_mean[cv_results$best_iteration])
}


set.seed(2025)
study <- optuna$create_study(direction = "minimize")
study$optimize(objective, n_trials = 50)

best_params <- study$best_params
best_params$objective <- "reg:squarederror"
print(best_params)

$eta [1] 0.2586295

$max_depth [1] 12

$min_child_weight [1] 4

$subsample [1] 0.9921435

$colsample_bytree [1] 0.5823858

$gamma [1] 0.08316905

$lambda [1] 4.260907

$alpha [1] 0.2049112

TRAIN/TEST SPLIT

# Set seed
set.seed(2025)

# Split by index
train_index <- createDataPartition(y, p = 0.8, list = FALSE)
X_train <- X[train_index, ]
y_train <- y[train_index]
X_test <- X[-train_index, ]
y_test <- y[-train_index]

# Convert to xgb.DMatrix
dtrain <- xgb.DMatrix(data = as.matrix(X_train), label = y_train)
dtest <- xgb.DMatrix(data = as.matrix(X_test), label = y_test)

TRAIN MODEL

# Set seed
set.seed(2025)

# Training with parallel processing
final_model <- xgb.train(
  params = list(
    eta = best_params$eta,
    max_depth = best_params$max_depth,
    gamma = best_params$gamma,
    colsample_bytree = best_params$colsample_bytree,
    min_child_weight = best_params$min_child_weight,
    subsample = best_params$subsample,
    objective = "reg:squarederror",
    eval_metric = "rmse"
  ),
  data = dtrain,
  nrounds = 1000,
  watchlist = list(train = dtrain, test = dtest),
  early_stopping_rounds = 20,
  verbose = 1,
  nthread = detectCores() - 1
)

SAVE MODEL

# Create directory if it doesn't exist
if (!dir.exists("../data/models")) {
  dir.create("../data/models", recursive = TRUE)
}

# Save the final XGBoost model
saveRDS(final_model, file = "../data/models/xgb_model.rds")

# Save the best parameters
saveRDS(best_params, file = "../data/models/xgb_best_params.rds")

cat("Model and parameters saved to ../data/models/")

# Load the final XGBoost model
final_model <- readRDS("../data/models/xgb_model.rds")

# Load the best parameters
best_params <- readRDS("../data/models/xgb_best_params.rds")

MODEL EVALUATION

set.seed(2025)

# Predict on test set
preds <- predict(final_model, as.matrix(X_test))

# Metrics 
rmse <- sqrt(mean((y_test - preds)^2))
mae <- mean(abs(y_test - preds))
r2 <- 1 - (sum((y_test - preds)^2) / sum((y_test - mean(y_test))^2))

cat("Model Evaluation Metrics:\n")

Model Evaluation Metrics:

cat("  RMSE:", rmse, "\n")

  RMSE: 0.3396146 

cat("  MAE :", mae, "\n")

  MAE : 0.2422538 

cat("  R²  :", r2, "\n\n")

  R²  : 0.425285 

# Residuals 
residuals <- y_test - preds
residual_df <- data.frame(
  actual = y_test,
  predicted = preds,
  residuals = residuals
)

# Predicted vs Actual
p1 <- residual_df %>%
  ggplot(aes(x = actual, y = predicted)) +
  geom_point(alpha = 0.5) +
  geom_abline(slope = 1, intercept = 0, color = "red") +
  theme_minimal() +
  labs(title = "Predicted vs Actual Crash Rates",
       x = "Actual",
       y = "Predicted")

# Residuals vs Predicted 
p2 <- residual_df %>%
  ggplot(aes(x = predicted, y = residuals)) +
  geom_point(alpha = 0.5, color = "blue") +
  geom_hline(yintercept = 0, linetype = "dashed", color = "red") +
  theme_minimal() +
  labs(title = "Residuals vs Predicted",
       x = "Predicted",
       y = "Residuals")

# Residual Density 
p3 <- ggplot(residual_df, aes(x = residuals)) +
    geom_histogram(
        aes(y = ..density..),  
        binwidth = 0.05, 
        fill = "steelblue", 
        color = "white",
        alpha = 0.7
    ) +
    geom_density(
        color = "red", 
        linewidth = 1, 
        adjust = 1.5  
    ) +
    theme_minimal() +
    labs(
        title = "Residual Distribution", 
        x = "Residuals", 
        y = "Density"   
    )


ggsave("../report/plots/predicted_vs_actual_values_plot.png", p1, width = 10, height = 6, dpi = 300)
ggsave("../report/plots/resisuals_vs_predicted_values_plot.png", p2, width = 10, height = 6, dpi = 300)
ggsave("../report/plots/residual_density_plot.png", p3, width = 10, height = 6, dpi = 300)

# Grid Plots
grid_plot <- p1 + p2 + p3 + plot_layout(ncol = 3)

ggsave(
  "../report/plots/residuals_diagnostic_grid.png",
  grid_plot,
  width = 18, 
  height = 6,
  dpi = 300
)

Model Evaluation Metrics: RMSE: 0.3396146 MAE : 0.2422538 R² : 0.425285 ## SHAP EXPLAINABILITY

# Compute SHAP values
shap_values <- shap.values(xgb_model = final_model, X_train = as.matrix(X_train))
shap_long <- shap.prep(shap_contrib = shap_values$shap_score, X_train = as.matrix(X_train))

# SHAP summary plot
print(shap.plot.summary(shap_long))

shap <- shap.plot.summary(shap_long)
ggsave("../report/plots/shap_summary_plot.png", shap, width = 10, height = 6, dpi = 300)

INSPECT TREE ENSEMBLE

xgb.plot.multi.trees(model = final_model)

# SHAP Dependence and Interaction Plots

# Interaction values
shap_interaction_values <- predict(
  final_model,
  as.matrix(X_train),
  predinteraction = TRUE
)

saveRDS(shap_interaction_values, "../data/models/shap_interactions.rds")

# Define the specific top features we want to analyze
selected_features <- c(
  "post_pandemic",
  "pct_age_65_plus",
  "pct_graduate_degree",
  "pct_in_labor_force",
  "median_income",
  "pct_commute_medium"
)

# Create pretty names for the selected features
feature_names <- c(
  "Pre vs Post Pandemic",
  "% Population 65+ Years",
  "% with Graduate Degree",
  "% in Labor Force",
  "Median Income",
  "% with Medium Commute"
)

# Verify all selected features exist in X_train
missing_features <- setdiff(selected_features, colnames(X_train))
if (length(missing_features) > 0) {
  stop(paste("The following features are missing from X_train:", 
             paste(missing_features, collapse = ", ")))
}

# Generate PDP plots with error handling
pdp_plots <- map2(selected_features, feature_names, function(f, name) {
  tryCatch({
    pd <- pdp::partial(
      final_model, 
      pred.var = f, 
      train = as.matrix(X_train), 
      grid.resolution = 30
    )
    
    autoplot(pd) +
      ggtitle(name) +
      theme_minimal() +
      theme(
        axis.title = element_blank(),
        axis.text = element_text(size = 8),
        plot.title = element_text(size = 10, hjust = 0.5),
        plot.margin = unit(c(1, 1, 1, 1), "lines")
      )
  }, error = function(e) {
    message(paste("Error generating PDP for", f, ":", e$message))
    return(NULL)
  })
})

# Remove any NULL entries from failed plots
pdp_plots <- purrr::compact(pdp_plots)

# Create directory if it doesn't exist
if (!dir.exists("../report/plots/pdp")) {
  dir.create("../report/plots/pdp", recursive = TRUE)
}

# Save individual plots
walk2(pdp_plots, selected_features, ~{
  if (!is.null(.x)) {
    ggsave(
      filename = paste0("../report/plots/pdp/pdp_", .y, ".png"),
      plot = .x,
      width = 5,
      height = 4,
      dpi = 300
    )
  }
})

# Create and save grid plot
if (length(pdp_plots) > 0) {
  grid_plot <- gridExtra::grid.arrange(
    grobs = pdp_plots, 
    ncol = 2,
    top = "Partial Dependence Plots for Key Features"
  )
  
  ggsave(
    filename = "../report/plots/pdp/pdp_grid.png",
    plot = grid_plot,
    width = 10,
    height = 12,
    dpi = 300
  )
  
  # Optionally display the grid plot
  print(grid_plot)
} else {
  warning("No PDP plots were successfully generated.")
}